import os.path
import cv2
from ultralytics import YOLO
import glob
import numpy as np
from PIL import Image
from ctypes import *
import ctypes
class YJ_OKNG_Classify():
    def __init__(self):
        self.C_model = YOLO(r'AI_ICCN/weight/best_YJ_OKNG_0112.pt')
    def run_img(self,image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        res = self.C_model(image)
        classes = res[0].names
        class_name = classes[res[0].probs.top1]
        conf = res[0].probs.top1conf
        return class_name,conf

class YJAC_OKNG_Classify():
    def __init__(self):
        self.AC_model = YOLO(r'AI_ICCN/weight/best_AC_0129.pt')
    def run_img(self,image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        res = self.AC_model(image)
        classes = res[0].names
        class_name = classes[res[0].probs.top1]
        conf = res[0].probs.top1conf
        return class_name,conf
def sort_points(points):
    # 定义一个小的阈值来比较浮点数，决定两个点是否在同一个组
    epsilon = 10

    def is_close(a, b, eps=epsilon):
        return abs(a - b) < eps

    # 按y坐标排序来找出上下组
    points.sort(key=lambda p: p[1])
    bottom_group = [p for p in points if is_close(p[1], points[0][1])]
    top_group = [p for p in points if is_close(p[1], points[-1][1])]

    # 移除上下组
    remaining_points = [p for p in points if p not in bottom_group and p not in top_group]
    # remaining_points  = points
    # 按x坐标排序来找出左右组
    remaining_points.sort(key=lambda p: p[0])
    left_group = [p for p in remaining_points if is_close(p[0], remaining_points[0][0])]
    right_group = [p for p in remaining_points if is_close(p[0], remaining_points[-1][0])]

    # 对每组内部进行排序
    bottom_group.sort(key=lambda p: (p[0], p[1]))
    top_group.sort(key=lambda p: (p[0], p[1]))
    left_group.sort(key=lambda p: (p[1], p[0]))  # 注意这里的顺序是先按y后按x
    right_group.sort(key=lambda p: (p[1], p[0]))  # 同上

    return bottom_group, top_group, left_group, right_group

#计算拟合直线的角度
def calc_angle(points):
    '''
    param: points e.g.
    [(1,2),(2,3),...]
    :return: angle
    '''
    if len(points) <= 1:
        return None
    points = np.array(points)
    output = cv2.fitLine(points, cv2.DIST_L2, 0, 1e-2, 1e-2)
    x = output[1] / output[0]
    # 判断斜率为正或负
    if x>=0:
        angle = np.rad2deg(np.arctan(x))
    else:
        angle = 180 + np.rad2deg(np.arctan(x))
    return angle

class YJ_Det():
    def __init__(self):
        self.model = YOLO(r'AI_ICCN/weight/best_YJ_Det_0112.pt')
        self.cls_model = YJ_OKNG_Classify()
        self.ac_cls_model = YJAC_OKNG_Classify()
        self.dll = ctypes.cdll.LoadLibrary(r"./dll/PTTLibrary.dll")
    def load_ptt(self,ptt_path):
        Width = ctypes.c_int(0)  # Image Width
        Length = ctypes.c_int(0)  # Image Length
        pixelSize = ctypes.c_int(0)  # Bytes of Pixel (Image)     3:BGR(Blue Green Red), 2:BR(Blue Red)
        ResolutionX = ctypes.c_float(0)  # Pixel Resoultion X(um)
        ResolutionY = ctypes.c_float(0)  # Pixel Resoultion Y(um)

        ROIStartX = ctypes.c_float(0)  # ROI StartX(Pixel)
        ROIStartY = ctypes.c_float(0)  # ROI StartY (Pixel)
        ROIWidth = ctypes.c_float(0)  # ROI Width (pixel)
        ROILength = ctypes.c_float(0)  # ROI Lengh (PIXEL)

        ptt = self.dll.ReadPTTBufferSizeExtend(ptt_path,
                                          ctypes.byref(Width),
                                          ctypes.byref(Length),
                                          ctypes.byref(pixelSize),
                                          ctypes.byref(ResolutionX),
                                          ctypes.byref(ResolutionY),
                                          ctypes.byref(ROIStartX),
                                          ctypes.byref(ROIStartY),
                                          ctypes.byref(ROIWidth),
                                          ctypes.byref(ROILength))
        heightMap = (c_float * (Width.value * Length.value))()
        IamgeMap = (c_ubyte * (Width.value * Length.value * pixelSize.value))()
        ptt2 = self.dll.ReadPTTBuffer(ptt_path, heightMap, IamgeMap)
        height_img = np.ctypeslib.as_array(heightMap)
        img_map = np.ctypeslib.as_array(IamgeMap)
        grayImage = height_img.reshape(Length.value, Width.value)
        img = img_map.reshape(Length.value, Width.value, pixelSize.value)
        # img = np.flip(img,axis=0)
        # grid = pv.ImageData(dimensions=(Width.value, Length.value, 1),spacing=(int(ResolutionX.value),int(ResolutionY.value),1.0))
        # grid.point_data.set_array(heightMap,'a')
        # grid.spacing
        # terrain = grid.warp_by_scalar(factor=1 / 20)
        # terrain.texture_map_to_plane(use_bounds=True, inplace=True)
        # texture = pv.Texture(img)
        # terrain.plot(texture=texture)
        return img, grayImage

    def _clc_mean(self,mean_list):
        mean = np.mean(mean_list)
        for x in mean_list:
            var =100*(np.abs(x - mean) / mean)
            # print(var * 100, '%')
            if var >5:
                print('Height NG!')
                return 'NG'
        return 'OK'
    def start_classify(self,datalist,ng_dict,ptt_root):
        res_dict = {}
        for data in datalist:
            key, image = data[0], data[2]
            img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            res = self.model(img)
            names = res[0].names
            h, w = res[0].orig_shape
            xyxys = boxes = res[0].boxes.xyxy
            xywh = res[0].boxes.xywh
            xywh_arr = np.asarray(xywh.cpu())
            xy_arr = xywh_arr[:, 0:2]
            xy_arr = [tuple(p) for p in xy_arr]  # lst2 = [(3, 3), (3, 2), (1, 1), (1, 2), (2, 1)]
            bottom, top, left, right = sort_points(xy_arr)

            # # 拟合直线计算角度
            # # print("Bottom group:", bottom)
            # # print("Top group:", top)
            # # print("Left group:", left)
            # # print("Right group:", right)
            # angle_b = calc_angle(bottom)
            # angle_t = calc_angle(top)
            # angle_l = calc_angle(left)
            # angle_r = calc_angle(right)
            # print(angle_b, angle_t, angle_l, angle_r)
            # if np.abs(angle_b[0]) > 1 or np.abs(angle_t[0]) > 1:
            #     res_dict[key] = '偏位'
            #     print('偏位')
            #     continue
            # elif angle_l != None and angle_r != None:
            #     if np.abs(angle_l - 90) > 1 or np.abs(angle_r - 90) > 1:
            #         res_dict[key] = '偏位'
            #         print('偏位')
            #         continue
            ori_img = image.copy()
            pttpath = os.path.join(ptt_root, key+'.ptt')
            ac_path = os.path.join(ptt_root,key+'_AC.jpg')
            if os.path.exists(ac_path):
                ac_img = cv2.imdecode(np.fromfile(ac_path,dtype= np.uint8),-1)
            else:
                ac_img = None
            res_yj = 'OK'
            if os.path.exists(pttpath):
                img_0, grayImage = self.load_ptt(pttpath)
                mean_list = []
                for i, xyxy in enumerate(xyxys):
                    H = xyxy[3] - xyxy[1]
                    W = xyxy[2] - xyxy[0]
                    c_x = (int(xyxy[1]) + int(xyxy[3])) / 2
                    c_y = (int(xyxy[0]) + int(xyxy[2])) / 2
                    x1_ = c_x - H

                    x2_ = c_x + H
                    y1_ = c_y - W
                    y2_ = c_y + W
                    x1_ptt = c_x - H / 4
                    x2_ptt = c_x + H / 4
                    y1_ptt = c_y - W / 4
                    y2_ptt = c_y + W / 4
                    crop = ori_img[int(x1_):int(x2_), int(y1_):int(y2_), :]
                    if ac_img is not None:
                        crop_ac = ac_img[int(x1_):int(x2_), int(y1_):int(y2_), :]
                        try:
                            res_ac = self.ac_cls_model.run_img(crop_ac)
                        except Exception as e:
                            print(e)
                            res_ac=None
                    else:
                        res_ac = None


                    crop_ptt = grayImage[max(0, int(x1_ptt)):int(x2_ptt), max(0, int(y1_ptt)):int(y2_ptt)]
                    # measure_area =
                    max_height = np.max(crop_ptt)
                    mean_list.append(max_height)
                    # 检测每个脚是否NG
                    try:
                        res = self.cls_model.run_img(crop)
                    except Exception as e:
                        print(e)
                        continue
                    if res_ac:
                        if res[0] == 'NG' or res_ac[0] == '0':
                            res_yj='NG'
                    else:
                        if res[0] == 'NG':
                            res_yj='NG'
                            # print('NG')
                            # continue

                if res_yj == 'NG':
                    ng_dict[key] = 'NG'
                if len(mean_list):
                    res = self._clc_mean(mean_list)
                    if res =='NG':
                        print(key,'.jpg: LIFT')
                        ng_dict[key] = 'L_LIFT'
                else:
                    ng_dict[key] = 'NG'

            else:
                for i, xyxy in enumerate(xyxys):
                    H = xyxy[3] - xyxy[1]
                    W = xyxy[2] - xyxy[0]
                    c_x = (int(xyxy[1]) + int(xyxy[3])) / 2
                    c_y = (int(xyxy[0]) + int(xyxy[2])) / 2
                    x1_ = c_x - H
                    x2_ = c_x + H
                    y1_ = c_y - W
                    y2_ = c_y + W
                    crop = ori_img[int(x1_):int(x2_), int(y1_):int(y2_), :]
                    # 检测每个脚是否NG
                    res = self.cls_model.run_img(crop)
                    if res[0] == 'NG':
                        res_yj = 'NG'
                        print('NG')
                        continue

                if res_yj == 'NG':
                    ng_dict[key] = 'NG'

        return ng_dict